import torch as pt
from python_ai.common.xcommon import *


class DepthSeparableConv(pt.nn.Module):

    def __init__(self, in_ch, out_ch, strides=(1, 1), **kwargs):
        super().__init__(**kwargs)

        self.depth = pt.nn.Conv2d(in_ch, in_ch, (3, 3), strides, 1, groups=in_ch, bias=False)
        self.bn1 = pt.nn.BatchNorm2d(in_ch)
        self.relu1 = pt.nn.ReLU()
        self.ptr = pt.nn.Conv2d(in_ch, out_ch, (1, 1), (1, 1), 0)
        self.bn2 = pt.nn.BatchNorm2d(out_ch)
        self.relu2 = pt.nn.ReLU()

    def forward(self, x):
        x = self.depth(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.ptr(x)
        x = self.bn2(x)
        x = self.relu2(x)
        return x


cfg = [
    (64, 1),
    (128, 2),
    (128, 1),
    (256, 2),
    (256, 1),
    (512, 2),
    (512, 1),
    (512, 1),
    (512, 1),
    (512, 1),
    (512, 1),
    (1024, 2),
    (1024, 1),
]


class MobileNet(pt.nn.Module):

    def __init__(self, in_ch, n_cls, **kwargs):
        super().__init__(**kwargs)

        start_ch = 32
        self.conv1 = pt.nn.Sequential(
            pt.nn.Conv2d(in_ch, start_ch, (3, 3), (2, 2), 1, bias=False),
            pt.nn.BatchNorm2d(start_ch),
            pt.nn.ReLU(),
        )

        self.main_layers = self.make_layers(32)

        self.avg = pt.nn.AdaptiveAvgPool2d((1, 1))

        self.fc = pt.nn.Linear(cfg[-1][0], n_cls)

    def make_layers(self, in_ch):
        layers = []
        for out_ch, stride in cfg:
            layers.append(DepthSeparableConv(in_ch, out_ch, (stride, stride)))
            in_ch = out_ch
        return pt.nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.main_layers(x)
        x = self.avg(x)
        x = pt.squeeze(x, dim=3)
        x = pt.squeeze(x, dim=2)
        x = self.fc(x)
        x = pt.nn.LogSoftmax(dim=1)(x)
        return x


if '__main__' == __name__:
    import python_ai.common.read_data.cat_dog_teachers as cat_dog
    import os
    import sys
    import python_ai.common.routine.pt_get_device as get_device
    import python_ai.common.routine.pt_compile_fit_evaluate as routine
    from torch.utils.data import TensorDataset, DataLoader
    import matplotlib.pyplot as plt
    import python_ai.common.routine.plot_loss_and_metric as plot

    pt.random.manual_seed(1)

    sep('Device')
    device = get_device.get_device()

    sep('Model')
    model = MobileNet(3, 2).to(device)
    print(model)

    sep('Test model')
    x = pt.zeros((4, 3, 224, 224), dtype=pt.float32).to(device)
    pred = model(x)
    print('pred', pred.size())

    sep('Cat and dog')

    VER = 'v6.0'
    IS_GO_ON = True
    ALPHA = 1e-6
    SIZE = 224
    BATCH_SIZE = 16
    N_EPOCHS = 10
    BASE_EPOCH = None
    BASE_DIR, FILE_NAME = os.path.split(__file__)
    dir = '../../../../../../large_data/DL1/_many_files/cats_and_dogs_filtered/train'
    IMG_DIR = BASE_DIR + '/' + dir
    LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)
    SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
    SAVE_PATH = os.path.join(SAVE_DIR, 'weight')

    (x_train, y_train), (x_val, y_val), (x_test, y_test) = cat_dog.load_data(IMG_DIR, 1, SIZE, SIZE, 'cats', 'dogs', ch_first=True)

    x_train = pt.Tensor(x_train)
    x_test = pt.Tensor(x_test)
    x_val = pt.Tensor(x_val)
    y_train = pt.Tensor(y_train)
    y_test = pt.Tensor(y_test)
    y_val = pt.Tensor(y_val)
    ds_train = TensorDataset(x_train, y_train)
    ds_test = TensorDataset(x_test, y_test)
    ds_val = TensorDataset(x_val, y_val)

    dl_train = DataLoader(ds_train, BATCH_SIZE, drop_last=True)
    dl_test = DataLoader(ds_test, BATCH_SIZE, drop_last=False)
    dl_val = DataLoader(ds_val, BATCH_SIZE, drop_last=False)

    model_dict = routine.compile(
        device,
        model,
        pt.nn.NLLLoss(),
        pt.optim.Adam,
        ALPHA,
        routine.accuracy,
    )

    need_train = False
    if os.path.exists(SAVE_DIR) and len(os.listdir(SAVE_DIR)) > 0:
        if BASE_EPOCH is None:
            BASE_EPOCH = routine.get_max_epoch(SAVE_PATH)
        load_path = SAVE_PATH + '.' + str(BASE_EPOCH)
        print(f'Loading {load_path} ...')
        sdict = pt.load(load_path)
        model.load_state_dict(sdict)
        print('Loaded.')
        if IS_GO_ON:
            need_train = True
    if need_train:
        loss_his, acc_his, loss_his_val, acc_his_val = routine.fit(
            model_dict, dl_train, dl_val, N_EPOCHS, LOG_DIR, SAVE_PATH, 2, BASE_EPOCH
        )

    routine.evaluate(model_dict, dl_test)

    if need_train:
        plt_dict = dict(spr=1, spc=2, spn=0)
        plt.figure(figsize=[12, 6])
        plot.plot_loss_and_metric(plt_dict, loss_his, acc_his, loss_his_val, acc_his_val, BASE_EPOCH)
        plt.show()
